import sys
sys.path.append('..')
from site_packages.utils.job import DataOp
from models.traditional import Tfidf, WordEmbedding


def train():
    # 加载数据集
    data_train = DataOp.load_data('data_train')
    corpus = data_train['question1_clean'].tolist() + data_train['question2_clean'].tolist()

    # 训练tfidf
    tfidf = Tfidf()
    tfidf.train(corpus)

    # 训练word2vec
    EMBEDDING_PARAMS = {
        'min_count': 2,
        'size': 200,
        'workers': 4,
        'iter': 15,
    }
    word_embedding = WordEmbedding(model_type='word2vec', model_name='word2vec_wld', mode='train', **EMBEDDING_PARAMS)
    word_embedding.train(corpus)

    # 训练gensim fasttext
    FASTTEXT_PARAMS = {
        'window': 3,
        'min_count': 2,
        'size': 200,
        'workers': 4,
        'iter': 15,
        'word_ngrams':2,
    }

    fasttext_embedding = WordEmbedding(model_type='fasttext', model_name='fasttext_wld', mode='train', **FASTTEXT_PARAMS)
    fasttext_embedding.train(corpus)


if __name__ == '__main__':
    train()
